import torch
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
from utils import Singleton
def star2text_zh(star):
    if star == "star 1":
        return "消极的"
    elif star == "star 2":
        return "半消极的"
    elif star == "star 3":
        return "中立的"
    elif star == "star 4":
        return "半积极的"
    elif star == "star 5":
        return "积极的"
    else:
        return star


class SentimentAnalysis(metaclass=Singleton):
    def __init__(self):
        model_name=r"E:\chinese_sentiment\chinese_sentiment"
        model = AutoModelForSequenceClassification.from_pretrained(model_name)
        tokenizer = AutoTokenizer.from_pretrained(model_name,)
        self.classifier = pipeline(
            "sentiment-analysis", model=model, tokenizer=tokenizer,
            device="cuda:0" if torch.cuda.is_available() else "cpu"
        )

    def sentiment_analysis_single(self, text: str):
        label = ''
        score = 0
        try:
            prediction = self.classifier(text)
            label = star2text_zh(prediction[0]['label'])
            score = prediction[0]['score']
        except Exception as e:
            print("Error in index:", e)
            print(f"fault data: {text}")
        return label, score